#!/usr/bin/env python3
"""
OpenMPI MCP Server

面向 openmpi 软件的 MCP server，提供 MPI 程序编译、执行和管理的命令行工具封装。

MCP tools 列表：
1. openmpi_compile - 编译 C 程序为 MPI 可执行文件
2. openmpi_run_basic - 在单节点上运行 MPI 程序
3. openmpi_run_cluster - 在多节点集群环境中运行 MPI 程序
4. openmpi_info - 获取 OpenMPI 版本和配置信息
5. openmpi_server - 启动和管理 OpenMPI 服务器
6. openmpi_clean - 清理 MPI 会话
7. openmpi_dvm - 启动动态虚拟机
8. openmpi_orte_commands - 执行 ORTE 相关命令

所有工具返回统一 JSON 结构：
{
  "success": bool,
  "command": str,
  "exit_code": int,
  "stdout": str,
  "stderr": str
}
"""

import subprocess
import json
from typing import List, Optional
from mcp.server.fastmcp import FastMCP, Context

mcp = FastMCP("OpenMPI MCP Server")

@mcp.tool()
def openmpi_compile(source_file: str, output_file: str, working_dir: str = ".") -> dict:
    """
    编译 C 程序为 MPI 可执行文件
    
    参数:
        source_file: C 源文件路径
        output_file: 输出可执行文件路径
        working_dir: 工作目录，默认为当前目录
        
    返回:
        JSON 结构包含执行结果
    """
    try:
        cmd = ["mpicc", source_file, "-o", output_file]
        result = subprocess.run(cmd, cwd=working_dir, capture_output=True, text=True, timeout=30)
        
        return {
            "success": result.returncode == 0,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"mpicc {source_file} -o {output_file}",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def openmpi_run_basic(executable: str, num_processes: int = 1, allow_run_as_root: bool = True, working_dir: str = ".") -> dict:
    """
    在单节点上运行 MPI 程序
    
    参数:
        executable: 可执行文件路径
        num_processes: 进程数量，默认为1
        allow_run_as_root: 是否允许以root运行，默认为True
        working_dir: 工作目录，默认为当前目录
        
    返回:
        JSON 结构包含执行结果
    """
    try:
        cmd = ["mpirun"]
        if allow_run_as_root:
            cmd.append("--allow-run-as-root")
        cmd.extend(["-np", str(num_processes), executable])
        
        result = subprocess.run(cmd, cwd=working_dir, capture_output=True, text=True, timeout=30)
        
        return {
            "success": result.returncode == 0,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"mpirun -np {num_processes} {executable}",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def openmpi_run_cluster(executable: str, hosts: List[str] = None, hostfile: str = None, num_processes: int = 1, 
                       oversubscribe: bool = False, working_dir: str = ".", rankfile: str = None, 
                       prefix: str = None, report_bindings: bool = False) -> dict:
    """
    在多节点集群环境中运行 MPI 程序
    
    参数:
        executable: 可执行文件路径
        hosts: 主机列表，如 ["node1", "node2"]
        hostfile: hostfile文件路径
        num_processes: 进程数量
        oversubscribe: 是否允许超订阅，默认为False
        working_dir: 工作目录
        rankfile: rankfile文件路径
        prefix: MPI安装前缀路径
        report_bindings: 是否报告进程绑定信息
        
    返回:
        JSON 结构包含执行结果
    """
    try:
        cmd = ["mpirun", "--allow-run-as-root"]
        
        if oversubscribe:
            cmd.append("--oversubscribe")
            
        if hosts:
            cmd.extend(["-H", ",".join(hosts)])
            
        if hostfile:
            cmd.extend(["--hostfile", hostfile])
            
        if rankfile:
            cmd.extend(["-rf", rankfile])
            
        if prefix:
            cmd.extend(["--prefix", prefix])
            
        if report_bindings:
            cmd.append("--report-bindings")
            
        cmd.extend(["-np", str(num_processes)])
        
        if working_dir and working_dir != ".":
            cmd.extend(["--wdir", working_dir])
            
        cmd.append(executable)
        
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
        
        return {
            "success": result.returncode == 0,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"mpirun cluster execution",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def openmpi_info(version: bool = False, help: bool = False) -> dict:
    """
    获取 OpenMPI 版本和配置信息
    
    参数:
        version: 是否只显示版本信息
        help: 是否显示帮助信息
        
    返回:
        JSON 结构包含执行结果
    """
    try:
        cmd = ["ompi_info"]
        
        if version:
            cmd.append("-V")
        elif help:
            cmd.append("-h")
            
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        
        return {
            "success": result.returncode == 0,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": "ompi_info",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def openmpi_server(daemon: bool = False, report_uri: bool = False) -> dict:
    """
    启动和管理 OpenMPI 服务器
    
    参数:
        daemon: 是否以守护进程方式运行
        report_uri: 是否报告服务器URI
        
    返回:
        JSON 结构包含执行结果
    """
    try:
        cmd = ["ompi-server"]
        
        if daemon:
            cmd.append("-d")
        elif report_uri:
            cmd.append("-r")
            cmd.append("server_uri")
            
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        
        return {
            "success": result.returncode == 0,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": "ompi-server",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def openmpi_clean(verbose: bool = False, debug: bool = False, help: bool = False) -> dict:
    """
    清理 MPI 会话
    
    参数:
        verbose: 是否显示详细输出
        debug: 是否显示调试信息
        help: 是否显示帮助信息
        
    返回:
        JSON 结构包含执行结果
    """
    try:
        cmd = ["ompi-clean"]
        
        if verbose:
            cmd.append("-v")
        elif debug:
            cmd.append("-d")
        elif help:
            cmd.append("-h")
            
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        
        return {
            "success": result.returncode == 0,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": "ompi-clean",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def openmpi_dvm(allow_run_as_root: bool = True, report_uri: bool = False) -> dict:
    """
    启动动态虚拟机
    
    参数:
        allow_run_as_root: 是否允许以root运行
        report_uri: 是否报告DVM URI
        
    返回:
        JSON 结构包含执行结果
    """
    try:
        cmd = ["ompi-dvm"]
        
        if allow_run_as_root:
            cmd.append("--allow-run-as-root")
        if report_uri:
            cmd.append("-report-uri")
            cmd.append("dvm_uri")
            
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        
        return {
            "success": result.returncode == 0,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": "ompi-dvm",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def openmpi_orte_commands(command: str, help: bool = False, version: bool = False, allow_run_as_root: bool = True) -> dict:
    """
    执行 ORTE 相关命令
    
    参数:
        command: ORTE命令名称 (orte-clean, orte-server, orted, orterun)
        help: 是否显示帮助信息
        version: 是否显示版本信息
        allow_run_as_root: 是否允许以root运行（仅orterun）
        
    返回:
        JSON 结构包含执行结果
    """
    try:
        if command not in ["orte-clean", "orte-server", "orted", "orterun"]:
            return {
                "success": False,
                "command": command,
                "exit_code": -1,
                "stdout": "",
                "stderr": f"Invalid ORTE command: {command}"
            }
            
        cmd = [command]
        
        if help:
            cmd.append("-h")
        elif version and command == "orterun":
            cmd.append("-V")
            
        if command == "orterun" and allow_run_as_root:
            cmd.append("--allow-run-as-root")
            
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        
        return {
            "success": result.returncode == 0,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

if __name__ == "__main__":
    mcp.run()